package com.liuhengjia.langchain4japply.api;

import com.liuhengjia.langchain4japply.config.assistant.Assistant;
import com.liuhengjia.langchain4japply.util.Jsons;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

/**
 * 函数调用 API
 */
@RestController
@RequestMapping("/api/function")
@RequiredArgsConstructor
public class FunctionAPI {
    private final ChatLanguageModel chatLanguageModel;
    private final Assistant functionAssistant;

    /**
     * 低级聊天
     *
     * @param message 用户信息
     * @return 模型响应
     */
    @GetMapping("/low/chat")
    public String lowChat(@RequestParam(value = "message") String message) {
        ToolSpecification toolSpecification = ToolSpecification.builder()
                .name("Calculator")
                .description("输入两个数，对这两个数求和")
                .parameters(JsonObjectSchema.builder()
                        .addIntegerProperty("a", "第一个参数")
                        .addIntegerProperty("b", "第二个参数")
                        .required("a", "b")
                        .build())
                .build();
        ChatResponse chatResponse = chatLanguageModel.doChat(
                ChatRequest.builder()
                        .messages(List.of(UserMessage.from(message)))
                        .parameters(
                                ChatRequestParameters.builder()
                                        .toolSpecifications(toolSpecification)
                                        .build())
                        .build());
        chatResponse.aiMessage().toolExecutionRequests().forEach(toolExecutionRequest -> {
            try {
                Class<?> aClass = Class.forName("com.liuhengjia.langchain4japply.tools." + toolExecutionRequest.name());
                Runnable runnable = (Runnable) Jsons.toJsonObject(toolExecutionRequest.arguments(), aClass);
                runnable.run();
            } catch (ClassNotFoundException e) {
                throw new RuntimeException(e);
            }
        });
        return chatResponse.aiMessage().text();
    }

    /**
     * 高级聊天
     *
     * @param message 用户信息
     * @return 模型响应
     */
    @GetMapping("/high/chat")
    public String highChat(@RequestParam(value = "message") String message) {
        return functionAssistant.chat(message);
    }
}
